from googleapiclient.discovery import build
from googleapiclient.errors import HttpError
import pandas as pd
from datetime import datetime
import os
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()
API_KEY = os.getenv('YOUTUBE_API_KEY')

# 设置系统代理
os.environ['HTTP_PROXY'] = 'http://127.0.0.1:7890'
os.environ['HTTPS_PROXY'] = 'http://127.0.0.1:7890'

def create_youtube_service():
    """创建YouTube API服务实例"""
    return build('youtube', 'v3', developerKey=API_KEY)

def search_videos_by_tag(youtube, tag, max_results=10):
    """根据标签搜索视频"""
    try:
        print(f"正在搜索标签: {tag}")
        search_response = youtube.search().list(
            q=tag,
            part='id,snippet',
            maxResults=max_results,
            type='video'
        ).execute()
        
        channel_ids = list(set([
            item['snippet']['channelId'] 
            for item in search_response['items']
        ]))
        
        print(f"找到 {len(channel_ids)} 个唯一频道")
        return channel_ids
    
    except HttpError as e:
        print(f'搜索视频时发生错误: {e}')
        return []

def get_channel_stats(youtube, channel_ids):
    """获取频道统计信息"""
    channel_data = []
    total = len(channel_ids)
    
    print(f"开始获取 {total} 个频道的统计信息...")
    
    for i, channel_id in enumerate(channel_ids, 1):
        try:
            print(f"处理第 {i}/{total} 个频道: {channel_id}")
            channel_response = youtube.channels().list(
                part='snippet,statistics',
                id=channel_id
            ).execute()
            
            if channel_response['items']:
                channel_info = channel_response['items'][0]
                stats = {
                    '频道名': channel_info['snippet']['title'],
                    '频道ID': channel_id,
                    '注册时间': channel_info['snippet']['publishedAt'],
                    '订阅数': channel_info['statistics']['subscriberCount'],
                    '总观看数': channel_info['statistics']['viewCount'],
                    '视频数量': channel_info['statistics']['videoCount']
                }
                channel_data.append(stats)
                print(f"成功获取频道信息: {stats['频道名']}")
                
        except HttpError as e:
            print(f'获取频道 {channel_id} 信息时发生错误: {e}')
            continue
    
    return channel_data

def test_api_connection():
    """测试API连接是否正常"""
    try:
        youtube = create_youtube_service()
        # 尝试一个简单的API调用
        request = youtube.channels().list(
            part="snippet",
            id="UC_x5XG1OV2P6uZZ5FSM9Ttw"  # Google Developers channel
        )
        response = request.execute()
        print("API连接测试成功！")
        return True
    except HttpError as e:
        print(f"API连接测试失败: {e}")
        return False

def main():
    """主函数"""
    # 创建YouTube服务实例
    youtube = create_youtube_service()
    
    # 设置要搜索的标签
    tag = input("请输入要搜索的标签: ")
    
    # 搜索视频并获取频道ID
    channel_ids = search_videos_by_tag(youtube, tag)
    
    # 获取频道统计信息
    channel_data = get_channel_stats(youtube, channel_ids)
    
    # 将数据转换为DataFrame并保存
    if channel_data:
        df = pd.DataFrame(channel_data)
        
        # 转换时间格式
        df['注册时间'] = pd.to_datetime(df['注册时间'], format='mixed').dt.strftime('%Y-%m-%d')
        
        # 保存为CSV文件
        filename = f'youtube_channels_{tag}_{datetime.now().strftime("%Y%m%d")}.csv'
        df.to_csv(filename, index=False, encoding='utf-8-sig')
        print(f'数据已保存到 {filename}')
    else:
        print('未找到相关数据')

if __name__ == '__main__':
    main() 